package com.sunchangpeng.command;

import java.io.*;
import java.util.*;

public class CommandLine {
    public static final String ERROR_PREFIX = "err> ";
    public static final String OUTPUT_PREFIX = "out> ";
    private static final String EMPTY = "";

    protected final List<String> cmdLine = new ArrayList<>();
    protected Map<String, String> env = null;
    protected boolean cleanEnvironment = false;
    protected File workingDirectory;
    protected String outPrefix = OUTPUT_PREFIX;
    protected String errPrefix = ERROR_PREFIX;

    protected OutputStream out = System.out;
    protected OutputStream err = System.err;

    protected CommandLine(final String command) {
        cmdLine.add(command);
    }

    public static CommandLine cmd(final String command) {
        return new CommandLine(command);
    }

    // ---------------------------------------------------------------- arguments
    public CommandLine workingDirectory(final File workDirectory) {
        this.workingDirectory = workDirectory;
        return this;
    }

    public CommandLine workingDirectory(final String workDirectory) {
        this.workingDirectory = new File(workDirectory);
        return this;
    }

    public CommandLine arg(final String argument) {
        cmdLine.add(argument);
        return this;
    }

    public CommandLine args(final String... arguments) {
        if (arguments != null && arguments.length > 0) {
            Collections.addAll(cmdLine, arguments);
        }
        return this;
    }

    public CommandLine outPrefix(final String prefix) {
        this.outPrefix = prefix;
        return this;
    }

    public CommandLine errPrefix(final String prefix) {
        this.errPrefix = prefix;
        return this;
    }

    public CommandLine out(final OutputStream out) {
        this.out = out;
        return this;
    }

    public CommandLine err(final OutputStream err) {
        this.err = err;
        return this;
    }

    public CommandLine env(final String key, final String value) {
        if (env == null) {
            env = new HashMap<>();
        }
        env.put(key, value);
        return this;
    }

    /**
     * When set to {@code true}, environment will not be copied from the
     * parent process and will be completly empty.
     */
    public CommandLine newEnv(final boolean clean) {
        cleanEnvironment = clean;
        return this;
    }

    public String commandLineText() {
        return join(cmdLine, ' ');
    }

    // ---------------------------------------------------------------- execute

    /**
     * Runs command and returns process result.
     */
    public ProcessResult run() {
        ByteArrayOutputStream baos = new ByteArrayOutputStream();
        out = err = baos;

        ProcessBuilder processBuilder = new ProcessBuilder();
        processBuilder.command(cmdLine);

        if (cleanEnvironment) {
            processBuilder.environment().clear();
        }

        if (env != null) {
            processBuilder.environment().putAll(env);
        }

        processBuilder.directory(workingDirectory);

        Process process = null;
        try {
            process = processBuilder.start();
        } catch (IOException ioex) {
            return writeException(baos, ioex);
        }

        StreamGobbler outputGobbler = new StreamGobbler(process.getInputStream(), out, outPrefix);
        StreamGobbler errorGobbler = new StreamGobbler(process.getErrorStream(), err, errPrefix);

        outputGobbler.start();
        errorGobbler.start();

        int result;

        try {
            result = process.waitFor();
        } catch (InterruptedException iex) {
            return writeException(baos, iex);
        }

        outputGobbler.waitFor();
        errorGobbler.waitFor();

        return new ProcessResult(result, baos.toString());
    }

    private ProcessResult writeException(final ByteArrayOutputStream baos, final Exception ex) {
        try {
            baos.write(errPrefix.getBytes());
        } catch (IOException ignore) {
        }

        ex.printStackTrace(new PrintStream(baos));
        return new ProcessResult(-1, baos.toString());
    }

    private String join(final Collection collection, final char separator) {
        if (collection == null) {
            return null;
        }

        if (collection.size() == 0) {
            return EMPTY;
        }

        final StringBuilder sb = new StringBuilder(collection.size() * 16);
        final Iterator it = collection.iterator();

        for (int i = 0; i < collection.size(); i++) {
            if (i > 0) {
                sb.append(separator);
            }

            sb.append(it.next());
        }

        return sb.toString();
    }

    public static class ProcessResult {
        public static final int OK = 0;

        private final int exitCode;
        private final String output;

        protected ProcessResult(final int existCode, final String output) {
            this.exitCode = existCode;
            this.output = output;
        }

        public boolean isSuccess() {
            return exitCode == OK;
        }

        /**
         * Returns process exit code.
         */
        public int getExitCode() {
            return exitCode;
        }

        /**
         * Returns process output.
         */
        public String getOutput() {
            return output;
        }
    }
}
